#!/usr/bin/env python
# -*- coding:utf-8 -*-
# @FileName  :async_mongoDB.py
# @Time      :2024/11/8 21:56
# @Author    :zhou_yi

""" 常用查询
$lt小于{'age': {'$lt': 20}}
$gt大于{'age': {'$gt': 20}}
$lte小于等于{'age': {'$lte': 20}}
$gte大于等于{'age': {'$gte': 20}}
$ne不等于{'age': {'$ne': 20}}
$in在范围内{'age': {'$in': [20, 23]}}
$nin不在范围内{'age': {'$nin': [20, 23]}}

$regex匹配正则{'name': {'$regex': '^M.*'}}name以M开头
$exists属性是否存在{'name': {'$exists': True}}name属性存在
$type类型判断{'age': {'$type': 'int'}}age的类型为int
$mod数字模操作{'age': {'$mod': [5, 0]}}年龄模5余0
$text文本查询{'$text': {'$search': 'Mike'}}text类型的属性中包含Mike字符串
$where高级条件查询{'$where': 'obj.fans_count == obj.follows_count'}自身粉丝数等于关注数

"""
from io import BytesIO
import asyncio
from bson import ObjectId
from pymongo import AsyncMongoClient
from pymongo.asynchronous.cursor import AsyncCursor
from pymongo.asynchronous.database import AsyncDatabase

from gridfs import AsyncGridFS
from gridfs.asynchronous.grid_file import AsyncGridOut


class CURD:
    def __init__(self, database: AsyncDatabase):
        self.db = database

    # 多条记录查询 返回异步对象
    def fetch_many(self, collection_name: str, query: dict) -> AsyncCursor:
        _cursor = self.db[collection_name].find(query)
        return _cursor

    # 单条记录查询
    async def fetch_one(self, collection_name: str, query: dict) -> dict:
        _cursor = await self.db[collection_name].find_one(query)
        return _cursor

    # 单条记录修改
    async def update_one(self, collection_name: str, query: dict, data: dict) -> tuple[bool, int, any]:
        try:
            _cursor = await self.db[collection_name].update_one(query, {'$set': data})
            return _cursor.acknowledged, _cursor.modified_count, _cursor.upserted_id
        except Exception:
            raise "数据修改时候出错"

    # 多条记录修改
    async def update_many(self, collection_name, query, data) -> tuple[bool, int, list]:
        try:
            _cursor = await self.db[collection_name].update_many(query, {'$set': data})
            return _cursor.acknowledged, _cursor.modified_count, _cursor.upserted_id
        except Exception:
            raise "数据修改时候出错"

    # 单条记录插入
    async def insert_one(self, collection_name: str, data: dict) -> tuple[bool, any]:
        try:
            _cursor = await self.db[collection_name].insert_one(data)
            return _cursor.acknowledged, _cursor.inserted_id
        except Exception:

            raise "数据存储发生错误 可能主键重复"

    # 多条记录插入
    async def insert_many(self, collection_name: str, data: list[dict]) -> tuple[bool, list]:
        try:
            _cursor = await self.db[collection_name].insert_many(data)
            return _cursor.acknowledged, _cursor.inserted_ids
        except Exception:
            raise "数据存储发生错误 可能主键重复"

    # 删除单条数据
    async def delete_one(self, collection_name: str, query: dict) -> tuple[bool, int, any]:
        try:
            Id: dict = await self.db[collection_name].find_one(query, {'_id': 1})
            _cursor = await self.db[collection_name].delete_one(query)
            return _cursor.acknowledged, _cursor.deleted_count, Id['_id']
        except Exception:
            raise "数据删除发生错误"

    # 删除多条数据
    async def delete_many(self, collection_name: str, query: dict) -> tuple[bool, int, list]:
        try:
            _ids = [await itm['_id'] for itm in self.db[collection_name].find(query, {'id': 1})]
            _cursor = await self.db[collection_name].delete_many(query)
            return _cursor.acknowledged, _cursor.deleted_count, _ids
        except Exception:
            raise "数据删除发生错误"

    # 求集合中文档id号的最大值 并生成新值 该集合的id是命名的 (前缀-数字)的形式
    async def collection_max_id_and_new_id(self, collection_name, sep='-') -> tuple[str, str]:
        """
            求 集合中 id 最大值,并生成 新增值id 适合 id=p-0016
            在表单新增数据的时候 按规则 自动分配给一个id值
        :param collection_name:
        :param sep: 字母与数字的分割符
        :return: p-0016（原数据中最大的） , p-0017（新）
        """

        try:
            _cursor = self.db[collection_name].find({}, {'_id': 1})  # 只取id号
            _ids: list[dict] = await _cursor.to_list()

            _id_prefix: str = _ids[0]['_id'].split(sep)[0]  # 取前缀 "p"

            _id_lst: list[int] = [int(itm['_id'].split(sep)[1]) for itm in _ids]  # 取数字部分 "0016"
            # max_id = sorted(_id_lst)[-1]  # 排序
            max_id: int = max(_id_lst)  # 最大值
            current_id: str = f"{int(max_id) + 1 :0>4}"
            return f"{_id_prefix}-{max_id:0>4}", f"{_id_prefix}-{current_id}"
        except Exception:
            raise '数据编号不符合 前缀-数字 的格式 例如：p-1003'


class FILES:
    def __init__(self, database):
        self.db = database

    async def find_one(self, collection, query) -> AsyncGridOut:
        _cursor = await AsyncGridFS(self.db, collection).find_one(query)
        return _cursor

    async def find_many(self, collection, query) -> list[AsyncGridOut]:
        _curser = AsyncGridFS(self.db, collection).find(query)
        return await _curser.to_list()

    # 为保证文件名的唯一性 文件存储同名会覆盖
    async def put(self, collection: str, data: bytes, filename: str) -> ObjectId:
        """
            存储文件 同名文件会覆盖
        :param collection: 集合名
        :param data: 二进制数据流
        :param filename: 文件名
        :return: 存储文件的ID
        """

        _gfs = AsyncGridFS(self.db, collection)

        _exist = await _gfs.find_one({'filename': filename})
        if _exist:
            await _gfs.delete(_exist._id)

        _cursor = await _gfs.put(data, filename=filename)
        return _cursor

    async def delete(self, collection, file_id: ObjectId):
        await AsyncGridFS(self.db, collection).delete(file_id)


class GFS:
    def __init__(self, database: AsyncDatabase, collection):
        self.gfs = AsyncGridFS(database, collection=collection)  # 相当于主文件夹

    # 增加一个文件
    async def put_file_by_name(self, b: bytes, name: str) -> tuple:
        _id = await self.gfs.put(data=b, filename=name)
        return _id, name

    # 所有文件的名字
    async def show_names(self, query=None) -> list[str]:
        files = self.gfs.find(query)
        return [itm.filename for itm in files]

    # 根据名字获取文件
    async def get_file_by_name(self, name: str) -> BytesIO:
        file = await self.gfs.get_version(filename=name)
        b: bytes = await file.read()
        io_ = BytesIO(b)
        return io_


# client = AsyncMongoClient(host='49.232.233.253', port=27017)
client = AsyncMongoClient(host='127.0.0.1', port=27017)
work = client.work
fs = client.files  # 文件数据库, 每个集合对应 存储的是 一个文件夹
Work = CURD(database=work)
Files = FILES(database=fs)

Gfs_point = GFS(database=fs, collection='point')
Gfs_cwrs = GFS(database=fs, collection='cwrs')

if __name__ == '__main__':
    import asyncio
    asyncio.run(Work.insert_one('clear_waster_records', {'_id': '12036'}))


    ...
